#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yang jiaxin'

import json
import openpyxl
import os
import re

ILLEGAL_CHARACTERS_RE = re.compile(r'[\000-\010]|[\013-\014]|[\016-\037]')

# pattern_dict = {
#     'dia01': '临床-诊断',
#     'ind01': '临床-指标',
#     'ser01': '分子-血清',
#     'ris01': '临床-风险因素',
#     'vac01': '分子-疫苗',
#     'phe01': '临床-症状',
#     'tre01': '临床-治疗',
#     'bio01': '分子-生化反应',
#     'mAb01': '分子-单抗',
#     'ner01': 'NER审编',
#     'ind02a': '临床-指标',
#     'ind02': '临床-指标',
#     'bio02': '分子-生化反应',
#     'dia02': '临床-诊断',
#     'vac02': '分子-疫苗',
#     'ner01': 'NER审编'
# }

entity_dict = {
    'clinical index': 'Clinical Index',
    'drug': 'Drug',
    'exam': 'Exam',
    'pathology': 'Pathology',
    'phenotype': 'Phenotype',
    'trail': 'Trial',
    'treatment': 'Treatment'
}

entity_code_dict = {
    40145: 'Behavior',
    39993: 'Body Part',
    39997: 'Cell',
    40146: 'Event',
    40147: 'Microbe',
    40897: 'Microbe',
    40569: 'Pathway',
    40144: 'Population',
    40150: 'Substance',
    40286: 'Vaccine',
    40148: 'Exposure'
}


def judge_text(subject_text, objects_text):
    """判断头实体文本和尾实体文本是否相同"""
    if subject_text == objects_text:
        text_consistency = 1
        return text_consistency
    elif subject_text != objects_text:
        text_consistency = 0
        return text_consistency


title_col = ['triple_id', 'articleId', 'paper_id', 'group', 'patternId', 'pattern', 'subject_name',
             'subject_type', 'subject_subLabels', 'subject_labelID', 'subject_start',
             'subject_end', 'subject_text', 'subject_annotate', 'relation', 'objects_name', 'objects_type',
             'objects_subLabels', 'objects_labelID', 'objects_start', 'objects_end',
             'objects_text', 'objects_annotate', 'correct', 'pmid', 'text_consistency', 'filename']

workbook = openpyxl.Workbook()
worksheet = workbook.active
worksheet.append(title_col)


def get_jsontoexcel(filedir, filename):
    with open(f'{filedir}/{filename}', 'r', encoding='utf-8') as f:
        data = json.load(f)
    for record in data['data']:

        triple_id = record['id']
        print(triple_id)
        print(filename)
        articleId = record['articleId']
        pattern = filename[int(filename.index("(")) + 1:int(filename.index(")"))]
        paper_id = articleId
        pmid = articleId.split("-")[-1]
        # 特殊文本编号的纠正
        if pmid == "999_20230317":
            pmid = "17225705"
        group = record['group']
        patternId = record['patternId']
        filename = filename
        relation = record['relation'].strip()
        correct = record['correct']
        # 头文本是否三元组
        if isinstance(record['subject'], str):
            subject_type = None
            subject_subLabels = None
            subject_labelID = None
            subject_start = None
            subject_end = None
            subject_text = None
            subject_annotate = None
            subject_name = record['subject']
        else:
            subject_type = record['subject']['type']
            if subject_type in entity_dict:
                subject_type = entity_dict[subject_type]
            subject_start = record['subject']['start']
            subject_end = record['subject']['end']
            subject_text = record['subject']['text']
            if subject_text is not None:
                subject_text = subject_text.strip()
                subject_text = ILLEGAL_CHARACTERS_RE.sub(r'-', subject_text)
                subject_name = subject_text[subject_start:subject_end].strip()
            else:
                subject_name = None
            subject_labelID = record['subject']['labelId']
            subject_annotate = record['subject']['annotate'] if 'annotate' in record['subject'] else None
            subject_subLabels = record['subject']['subLabels'] if 'subLabels' in record['subject'] else None
            if subject_subLabels is not None:
                if subject_subLabels in entity_code_dict:
                    subject_type = entity_code_dict[subject_subLabels]

        # 尾实体是否三元组
        if isinstance(record['objects'], str):
            objects_type = None
            objects_subLabels = None
            objects_labelID = None
            objects_start = None
            objects_end = None
            objects_text = None
            objects_annotate = None
            objects_name = record['objects']
        else:
            objects_type = record['objects']['type']
            if objects_type in entity_dict:
                objects_type = entity_dict[objects_type]
            objects_start = record['objects']['start']
            objects_end = record['objects']['end']
            objects_text = record['objects']['text']
            if objects_text is not None:
                objects_text = objects_text.strip()
                objects_text = ILLEGAL_CHARACTERS_RE.sub(r'-', objects_text)
                objects_name = objects_text[objects_start:objects_end].strip()
            else:
                objects_name = None
            objects_labelID = record['objects']['labelId']
            objects_annotate = record['objects']['annotate'] if 'annotate' in record['objects'] else None
            objects_subLabels = record['objects']['subLabels'] if 'subLabels' in record['objects'] else None
            if objects_subLabels is not None:
                if objects_subLabels in entity_code_dict:
                    objects_type = entity_code_dict[objects_subLabels]

        text_consistency = judge_text(subject_text, objects_text)

        insertdata = [triple_id, articleId, paper_id, group, patternId, pattern, subject_name, subject_type,
                      subject_subLabels, subject_labelID, subject_start, subject_end, subject_text, subject_annotate,
                      relation, objects_name, objects_type, objects_subLabels, objects_labelID, objects_start,
                      objects_end, objects_text, objects_annotate, correct, pmid, text_consistency, filename]
        print(insertdata)
        worksheet.append(insertdata)


folder_path = r'/Users/genesis/Seafile/Work/生物岛知识图谱/标注数据/肺癌/0415_batch0/'

for filename in os.listdir(folder_path):
    get_jsontoexcel(folder_path, filename)
workbook.save('/Users/genesis/Seafile/Work/生物岛知识图谱/标注数据/肺癌/LungCancer_0428_batch0.xlsx')
